{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from scipy import signal\n",
    "import mne\n",
    "from mne.preprocessing import ICA\n",
    "\n",
    "data_folder = '' # define the path of the data folder\n",
    "data_files = glob.glob(os.path.join(data_folder, '*.csv'))\n",
    "\n",
    "for data_file in data_files:\n",
    "    # load EEG data\n",
    "    data = pd.read_csv(data_file, header=0)\n",
    "    eeg_channels = data.columns[1:15]\n",
    "    eeg_data = data[eeg_channels].values\n",
    "    sfreq = 128  # sampling frequency is 128 Hz\n",
    "\n",
    "    # crate MNE raw object\n",
    "    ch_names = list(eeg_channels)\n",
    "    ch_types = ['eeg'] * len(ch_names)\n",
    "    info = mne.create_info(ch_names=ch_names, ch_types=ch_types, sfreq=sfreq)\n",
    "    raw = mne.io.RawArray(eeg_data.T, info=info)\n",
    "\n",
    "    # bandpass filter\n",
    "    raw.filter(l_freq=1.0, h_freq=50.0)\n",
    "\n",
    "    # create ICA object and fit it to raw data\n",
    "    ica = ICA(n_components=len(eeg_channels), random_state=0, max_iter=1000)  # 调整参数\n",
    "    ica.fit(raw)\n",
    "\n",
    "    # apply ICA to raw data\n",
    "    ica.exclude = []  \n",
    "    ica.apply(raw)\n",
    "\n",
    "    # get filtered data\n",
    "    filtered_data = raw.get_data()\n",
    "\n",
    "    # define frequency bands\n",
    "    freq_bands = {\n",
    "                # 'delta': (1, 4),\n",
    "                'theta': (4, 8),\n",
    "                'alpha': (8, 14),\n",
    "                'beta': (14, 31),\n",
    "                # 'gamma': (31, 50),\n",
    "                }\n",
    "\n",
    "    # parameters for PSD extraction\n",
    "    window_size = int(sfreq * 0.25)  # window size is 0.25 seconds\n",
    "    overlap = 0.0  # overlap between consecutive windows is 0.0\n",
    "    hop_size = int(window_size * (1 - overlap))\n",
    "\n",
    "    # one sample is 3 seconds data\n",
    "    segment_duration = 3 * sfreq \n",
    "    num_segments = int(np.floor(filtered_data.shape[1] / segment_duration))\n",
    "\n",
    "    num_channels = len(eeg_channels)\n",
    "    num_windows_per_segment = int(np.floor((segment_duration - window_size) / hop_size)) + 1\n",
    "    num_total_windows = num_windows_per_segment * num_segments\n",
    "    print(num_total_windows)\n",
    "\n",
    "    num_freq_bands = len(freq_bands)\n",
    "    psd_results = np.zeros((num_total_windows, num_channels * num_freq_bands))\n",
    "\n",
    "    # extract PSD features\n",
    "    window_idx = 0\n",
    "    for segment_idx in range(num_segments):\n",
    "        start = segment_idx * segment_duration\n",
    "        end = start + segment_duration\n",
    "        segment_data = filtered_data[:, start:end]\n",
    "\n",
    "        for window_start in range(0, segment_duration - window_size + 1, hop_size):\n",
    "            window_end = window_start + window_size\n",
    "            eeg_window = segment_data[:, window_start:window_end]\n",
    "            # welch PSD\n",
    "            frequencies, psd = signal.welch(eeg_window, fs=sfreq, nperseg=window_size)\n",
    "            for channel_idx in range(num_channels):\n",
    "                for band_idx, (band_name, (f_low, f_high)) in enumerate(freq_bands.items()):\n",
    "                    band_mask = np.logical_and(frequencies >= f_low, frequencies <= f_high)\n",
    "                    band_psd = psd[channel_idx, band_mask]\n",
    "                    PSDs[window_idx, channel_idx * num_freq_bands + band_idx] = np.sum(band_psd)\n",
    "                    PSDs = PSDs.reshape(num_total_windows, num_channels, num_freq_bands)\n",
    "            window_idx += 1\n",
    "\n",
    "    # PSDs : (num_total_windows, num_channels, num_freq_bands)\n",
    "    print(PSDs.shape)\n",
    "    "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "EEG",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.8.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
